{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 节块化，对各种变量进行初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class SubchannelBlock:\n",
    "    def __init__(self, n_channels, axial_nodes, ndx, ndxin):\n",
    "        self.n_channels = n_channels  # 燃料组件的子通道数\n",
    "        self.axial_nodes = axial_nodes  # 沿轴向的有效节点数\n",
    "        self.ndx = ndx  # 沿通道长度的有效控制容积数\n",
    "        self.ndxin = ndxin  # 沿输水管的有效控制容积数目\n",
    "        self.attributes = {}\n",
    "        self.initialize_attributes()\n",
    "    \n",
    "    def add_attribute(self, name):\n",
    "        \"\"\"添加新的属性（如液相含率、气相速度等）到子通道块\"\"\"\n",
    "        self.attributes[name] = np.zeros((self.n_channels, self.axial_nodes + 3))\n",
    "    \n",
    "    def set_attribute(self, name, values):\n",
    "        \"\"\"设置指定属性的初始值\"\"\"\n",
    "        if name in self.attributes:\n",
    "            self.attributes[name][:] = values\n",
    "        else:\n",
    "            raise ValueError(f\"Attribute {name} not defined in SubchannelBlock.\")\n",
    "    \n",
    "    def get_attribute(self, name):\n",
    "        \"\"\"获取指定属性的值\"\"\"\n",
    "        return self.attributes.get(name, None)\n",
    "    \n",
    "    def initialize_attributes(self):\n",
    "        \"\"\"初始化所有基础物理属性\"\"\"\n",
    "        # 添加基础属性\n",
    "        self.add_attribute(\"alpha_l\")            # 液相含率\n",
    "        self.add_attribute(\"alpha_g\")            # 气相含率\n",
    "        self.add_attribute(\"velocity_l\")         # 液相速度\n",
    "        self.add_attribute(\"velocity_g\")         # 气相速度\n",
    "        self.add_attribute(\"pressure\")           # 压力\n",
    "        self.add_attribute(\"enthalpy_l\")         # 液相焓\n",
    "        self.add_attribute(\"enthalpy_g\")         # 气相焓\n",
    "        self.add_attribute(\"mass\")               # 质量\n",
    "        self.add_attribute(\"density\")            # 混合物密度\n",
    "        self.add_attribute(\"flow\")               # 总流量\n",
    "        self.add_attribute(\"momentum\")           # 动量\n",
    "        self.add_attribute(\"mass_exchange\")      # 质量交换项\n",
    "        self.add_attribute(\"flow_l\")             # 液相流量\n",
    "        self.add_attribute(\"flow_v\")             # 气相流量\n",
    "        self.add_attribute(\"density_l\")          # 液相密度\n",
    "        self.add_attribute(\"density_v\")          # 气相密度\n",
    "        self.add_attribute(\"gravity_x\")          # 重力加速度\n",
    "        self.add_attribute(\"drag_coefficient\")   # 阻力系数\n",
    "        self.add_attribute(\"heat_flux_wall_liquid\")  # 墙面对液体的热流\n",
    "        self.add_attribute(\"heat_flux_wall_vapor\")  # 墙面对气体的热流\n",
    "        self.add_attribute(\"heat_flux_liquid\")   # 液相之间的热流\n",
    "        \n",
    "        # 添加中间变量\n",
    "        self.add_attribute(\"W_l_prime\")          # 液相质量交换速率的导数项\n",
    "        self.add_attribute(\"W_v_prime\")          # 气相质量交换速率的导数项\n",
    "        self.add_attribute(\"delta_U_l\")          # 液相速度变化量\n",
    "        self.add_attribute(\"delta_U_v\")          # 气相速度变化量\n",
    "        self.add_attribute(\"Psi_l_n\")            # 热流相关系数\n",
    "        self.add_attribute(\"delta_h_l\")          # 液相焓变化量\n",
    "        self.add_attribute(\"delta_h_v\")          # 气相焓变化量\n",
    "\n",
    "    def initialize_values(self):\n",
    "        \"\"\"初始化各属性的初始值\"\"\"\n",
    "        self.set_attribute(\"alpha_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.5))\n",
    "        self.set_attribute(\"alpha_g\", np.full((self.n_channels, self.axial_nodes + 3), 0.5))\n",
    "        self.set_attribute(\"velocity_l\", np.full((self.n_channels, self.axial_nodes + 3), 1.0))\n",
    "        self.set_attribute(\"velocity_g\", np.full((self.n_channels, self.axial_nodes + 3), 0.1))\n",
    "        self.set_attribute(\"pressure\", np.full((self.n_channels, self.axial_nodes + 3), 1e5))\n",
    "        self.set_attribute(\"enthalpy_l\", np.full((self.n_channels, self.axial_nodes + 3), 1e4))\n",
    "        self.set_attribute(\"enthalpy_g\", np.full((self.n_channels, self.axial_nodes + 3), 2e4))\n",
    "        self.set_attribute(\"mass\", np.full((self.n_channels, self.axial_nodes + 3), 1.0))\n",
    "        self.set_attribute(\"density\", np.full((self.n_channels, self.axial_nodes + 3), 1.2))\n",
    "        self.set_attribute(\"flow\", np.full((self.n_channels, self.axial_nodes + 3), 0.1))\n",
    "        self.set_attribute(\"momentum\", np.zeros((self.n_channels, self.axial_nodes + 3)))\n",
    "        self.set_attribute(\"mass_exchange\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"flow_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.05))\n",
    "        self.set_attribute(\"flow_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.05))\n",
    "        self.set_attribute(\"density_l\", np.full((self.n_channels, self.axial_nodes + 3), 1000.0))\n",
    "        self.set_attribute(\"density_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.6))\n",
    "        self.set_attribute(\"gravity_x\", np.full((self.n_channels, self.axial_nodes + 3), 9.81))\n",
    "        self.set_attribute(\"drag_coefficient\", np.full((self.n_channels, self.axial_nodes + 3), 0.1))\n",
    "        self.set_attribute(\"heat_flux_wall_liquid\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"heat_flux_wall_vapor\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"heat_flux_liquid\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        \n",
    "        # 中间变量初始化\n",
    "        self.set_attribute(\"W_l_prime\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"W_v_prime\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_U_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_U_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"Psi_l_n\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_h_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_h_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "\n",
    "# 示例使用\n",
    "n_channels = 7      # 子通道数\n",
    "axial_nodes = 10    # 轴向节点数\n",
    "ndx = 10            # 沿通道的有效控制容积数\n",
    "ndxin = 8           # 沿输水管的有效控制容积数目\n",
    "\n",
    "block = SubchannelBlock(n_channels, axial_nodes, ndx, ndxin)\n",
    "block.initialize_values()  # 初始化所有属性的默认值\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 加上了六守恒方程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class SubchannelBlock:\n",
    "    def __init__(self, n_channels, axial_nodes, ndx, ndxin):\n",
    "        self.n_channels = n_channels  # 燃料组件的子通道数\n",
    "        self.axial_nodes = axial_nodes  # 沿轴向的有效节点数\n",
    "        self.ndx = ndx  # 沿通道长度的有效控制容积数\n",
    "        self.ndxin = ndxin  # 沿输水管的有效控制容积数目\n",
    "        self.attributes = {}\n",
    "        self.initialize_attributes()\n",
    "    \n",
    "    def add_attribute(self, name):\n",
    "        \"\"\"添加新的属性（如液相含率、气相速度等）到子通道块\"\"\"\n",
    "        self.attributes[name] = np.zeros((self.n_channels, self.axial_nodes + 3))\n",
    "    \n",
    "    def set_attribute(self, name, values):\n",
    "        \"\"\"设置指定属性的初始值\"\"\"\n",
    "        if name in self.attributes:\n",
    "            self.attributes[name][:] = values\n",
    "        else:\n",
    "            raise ValueError(f\"Attribute {name} not defined in SubchannelBlock.\")\n",
    "    \n",
    "    def get_attribute(self, name):\n",
    "        \"\"\"获取指定属性的值\"\"\"\n",
    "        return self.attributes.get(name, None)\n",
    "    \n",
    "    def initialize_attributes(self):\n",
    "        \"\"\"初始化所有基础物理属性\"\"\"\n",
    "        self.add_attribute(\"alpha_l\")            # 液相含率\n",
    "        self.add_attribute(\"alpha_g\")            # 气相含率\n",
    "        self.add_attribute(\"velocity_l\")         # 液相速度\n",
    "        self.add_attribute(\"velocity_g\")         # 气相速度\n",
    "        self.add_attribute(\"pressure\")           # 压力\n",
    "        self.add_attribute(\"enthalpy_l\")         # 液相焓\n",
    "        self.add_attribute(\"enthalpy_g\")         # 气相焓\n",
    "        self.add_attribute(\"mass\")               # 质量\n",
    "        self.add_attribute(\"density\")            # 混合物密度\n",
    "        self.add_attribute(\"flow\")               # 总流量\n",
    "        self.add_attribute(\"momentum\")           # 动量\n",
    "        self.add_attribute(\"mass_exchange\")      # 质量交换项\n",
    "        self.add_attribute(\"flow_l\")             # 液相流量\n",
    "        self.add_attribute(\"flow_v\")             # 气相流量\n",
    "        self.add_attribute(\"density_l\")          # 液相密度\n",
    "        self.add_attribute(\"density_v\")          # 气相密度\n",
    "        self.add_attribute(\"gravity_x\")          # 重力加速度\n",
    "        self.add_attribute(\"drag_coefficient\")   # 阻力系数\n",
    "        self.add_attribute(\"heat_flux_wall_liquid\")  # 墙面对液体的热流\n",
    "        self.add_attribute(\"heat_flux_wall_vapor\")  # 墙面对气体的热流\n",
    "        self.add_attribute(\"heat_flux_liquid\")   # 液相之间的热流\n",
    "        \n",
    "        # 添加中间变量\n",
    "        self.add_attribute(\"W_l_prime\")          # 液相质量交换速率的导数项\n",
    "        self.add_attribute(\"W_v_prime\")          # 气相质量交换速率的导数项\n",
    "        self.add_attribute(\"delta_U_l\")          # 液相速度变化量\n",
    "        self.add_attribute(\"delta_U_v\")          # 气相速度变化量\n",
    "        self.add_attribute(\"Psi_l_n\")            # 热流相关系数\n",
    "        self.add_attribute(\"delta_h_l\")          # 液相焓变化量\n",
    "        self.add_attribute(\"delta_h_v\")          # 气相焓变化量\n",
    "\n",
    "    def initialize_values(self):\n",
    "        \"\"\"初始化各属性的初始值\"\"\"\n",
    "        self.set_attribute(\"alpha_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.5))\n",
    "        self.set_attribute(\"alpha_g\", np.full((self.n_channels, self.axial_nodes + 3), 0.5))\n",
    "        self.set_attribute(\"velocity_l\", np.full((self.n_channels, self.axial_nodes + 3), 1.0))\n",
    "        self.set_attribute(\"velocity_g\", np.full((self.n_channels, self.axial_nodes + 3), 0.1))\n",
    "        self.set_attribute(\"pressure\", np.full((self.n_channels, self.axial_nodes + 3), 1e5))\n",
    "        self.set_attribute(\"enthalpy_l\", np.full((self.n_channels, self.axial_nodes + 3), 1e4))\n",
    "        self.set_attribute(\"enthalpy_g\", np.full((self.n_channels, self.axial_nodes + 3), 2e4))\n",
    "        self.set_attribute(\"mass\", np.full((self.n_channels, self.axial_nodes + 3), 1.0))\n",
    "        self.set_attribute(\"density\", np.full((self.n_channels, self.axial_nodes + 3), 1.2))\n",
    "        self.set_attribute(\"flow\", np.full((self.n_channels, self.axial_nodes + 3), 0.1))\n",
    "        self.set_attribute(\"momentum\", np.zeros((self.n_channels, self.axial_nodes + 3)))\n",
    "        self.set_attribute(\"mass_exchange\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"flow_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.05))\n",
    "        self.set_attribute(\"flow_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.05))\n",
    "        self.set_attribute(\"density_l\", np.full((self.n_channels, self.axial_nodes + 3), 1000.0))\n",
    "        self.set_attribute(\"density_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.6))\n",
    "        self.set_attribute(\"gravity_x\", np.full((self.n_channels, self.axial_nodes + 3), 9.81))\n",
    "        self.set_attribute(\"drag_coefficient\", np.full((self.n_channels, self.axial_nodes + 3), 0.1))\n",
    "        self.set_attribute(\"heat_flux_wall_liquid\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"heat_flux_wall_vapor\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"heat_flux_liquid\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        \n",
    "        # 中间变量初始化\n",
    "        self.set_attribute(\"W_l_prime\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"W_v_prime\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_U_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_U_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"Psi_l_n\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_h_l\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "        self.set_attribute(\"delta_h_v\", np.full((self.n_channels, self.axial_nodes + 3), 0.0))\n",
    "    \n",
    "    def mass_conservation(self, dt, dx):\n",
    "        rho = self.get_attribute(\"density\")\n",
    "        F = self.get_attribute(\"flow\")\n",
    "        W = self.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        drho_dt = (rho[1:] - rho[:-1]) / dt\n",
    "        dF_dx = (F[1:] - F[:-1]) / dx\n",
    "        mass_balance = drho_dt + dF_dx + W[:-1]\n",
    "        return mass_balance\n",
    "    \n",
    "    def axial_momentum_conservation(self, dt, dx, A, K_U):\n",
    "        F = self.get_attribute(\"flow\")\n",
    "        P = self.get_attribute(\"pressure\")\n",
    "        rho = self.get_attribute(\"density\")\n",
    "        g_x = self.get_attribute(\"gravity_x\")\n",
    "        \n",
    "        F_l = self.get_attribute(\"flow_l\")\n",
    "        F_v = self.get_attribute(\"flow_v\")\n",
    "        U_l = self.get_attribute(\"velocity_l\")\n",
    "        U_v = self.get_attribute(\"velocity_g\")\n",
    "        \n",
    "        W_l = self.get_attribute(\"mass_exchange\")\n",
    "        W_l_prime = self.get_attribute(\"W_l_prime\")\n",
    "        delta_U_l = self.get_attribute(\"delta_U_l\")\n",
    "        W_v_prime = self.get_attribute(\"W_v_prime\")\n",
    "        delta_U_v = self.get_attribute(\"delta_U_v\")\n",
    "        \n",
    "        dF_dt = (F[1:] - F[:-1]) / dt\n",
    "        dF_dx = (F_l[1:] * U_l[1:] + F_v[1:] * U_v[1:] - F_l[:-1] * U_l[:-1] - F_v[:-1] * U_v[:-1]) / dx\n",
    "        dP_dx = (P[1:] - P[:-1]) / dx\n",
    "        \n",
    "        momentum_balance = (\n",
    "            dF_dt + dF_dx\n",
    "            + dP_dx * A\n",
    "            - W_l[:-1] * U_l[:-1]\n",
    "            - W_v_prime[:-1] * delta_U_v[:-1]\n",
    "            - K_U * np.abs(F[:-1]) * F[:-1]\n",
    "            + rho[:-1] * g_x[:-1]\n",
    "        )\n",
    "        return momentum_balance\n",
    "    \n",
    "    def transverse_momentum_conservation(self, dt, dx, S, K_v):\n",
    "        rho = self.get_attribute(\"density\")\n",
    "        V = self.get_attribute(\"velocity_g\")\n",
    "        F = self.get_attribute(\"flow\")\n",
    "        P = self.get_attribute(\"pressure\")\n",
    "        W = self.get_attribute(\"mass_exchange\")\n",
    "        W_v_prime = self.get_attribute(\"W_v_prime\")\n",
    "        delta_U_v = self.get_attribute(\"delta_U_v\")\n",
    "        g_y = 9.81  # 重力在y方向的加速度\n",
    "        \n",
    "        dV_dt = (V[1:] - V[:-1]) / dt\n",
    "        dFV_dx = (F[1:] * V[1:] - F[:-1] * V[:-1]) / dx\n",
    "        dF_dx = (F[1:] - F[:-1]) / dx\n",
    "        dP_dy = (P[1:] - P[:-1]) / dx  # 近似在x方向\n",
    "        \n",
    "        transverse_balance = (\n",
    "            9 * rho[:-1] * dV_dt + dFV_dx - V[:-1] * dF_dx\n",
    "            - 9 * dP_dy\n",
    "            + S * dx * K_v * W[:-1] * W[:-1]\n",
    "            + (W[:-1] * delta_U_v[:-1] + W_v_prime[:-1] * delta_U_v[:-1]) / (S * dx)\n",
    "            + 9 * rho[:-1] * g_y\n",
    "        )\n",
    "        return transverse_balance\n",
    "    \n",
    "    def enthalpy_conservation_liquid(self, dt, dx, A):\n",
    "        alpha = self.get_attribute(\"alpha_l\")\n",
    "        rho_l = self.get_attribute(\"density_l\")\n",
    "        h_l = self.get_attribute(\"enthalpy_l\")\n",
    "        F_l = self.get_attribute(\"flow_l\")\n",
    "        W_l = self.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        dh_l_dt = (h_l[1:] - h_l[:-1]) / dt\n",
    "        dFh_dx = (F_l[1:] * h_l[1:] - F_l[:-1] * h_l[:-1]) / dx\n",
    "        dF_dx = (F_l[1:] - F_l[:-1]) / dx\n",
    "        \n",
    "        enthalpy_balance_l = (\n",
    "            (1 - alpha[:-1]) * A * dx * rho_l[:-1] * dh_l_dt\n",
    "            + dx * dFh_dx - dx * h_l[:-1] * dF_dx + W_l[:-1] * h_l[:-1]\n",
    "        )\n",
    "        return enthalpy_balance_l\n",
    "    \n",
    "    def enthalpy_conservation_vapor(self, dt, dx, A):\n",
    "        alpha = self.get_attribute(\"alpha_g\")\n",
    "        rho_v = self.get_attribute(\"density_v\")\n",
    "        h_v = self.get_attribute(\"enthalpy_g\")\n",
    "        F_v = self.get_attribute(\"flow_v\")\n",
    "        W_v = self.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        dh_v_dt = (h_v[1:] - h_v[:-1]) / dt\n",
    "        dFh_dx = (F_v[1:] * h_v[1:] - F_v[:-1] * h_v[:-1]) / dx\n",
    "        dF_dx = (F_v[1:] - F_v[:-1]) / dx\n",
    "        \n",
    "        enthalpy_balance_v = (\n",
    "            alpha[:-1] * A * dx * rho_v[:-1] * dh_v_dt\n",
    "            + dx * dFh_dx - dx * h_v[:-1] * dF_dx + W_v[:-1] * h_v[:-1]\n",
    "        )\n",
    "        return enthalpy_balance_v\n",
    "    \n",
    "    def mixture_enthalpy_conservation(self, dt, dx, A):\n",
    "        rho = self.get_attribute(\"density\")\n",
    "        h = self.get_attribute(\"enthalpy_l\")  # 使用液相焓作为混合物焓\n",
    "        F_l = self.get_attribute(\"flow_l\")\n",
    "        h_l = self.get_attribute(\"enthalpy_l\")\n",
    "        F_v = self.get_attribute(\"flow_v\")\n",
    "        h_v = self.get_attribute(\"enthalpy_g\")\n",
    "        W_l = self.get_attribute(\"mass_exchange\")\n",
    "        W_v = self.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        dh_dt = (h[1:] - h[:-1]) / dt\n",
    "        dFh_dx = (F_l[1:] * h_l[1:] + F_v[1:] * h_v[1:] - F_l[:-1] * h_l[:-1] - F_v[:-1] * h_v[:-1]) / dx\n",
    "        dF_dx = (F_l[1:] + F_v[1:] - F_l[:-1] - F_v[:-1]) / dx\n",
    "        \n",
    "        enthalpy_balance_mixture = (\n",
    "            A * dx * rho[:-1] * dh_dt\n",
    "            + dx * dFh_dx - dx * h[:-1] * dF_dx\n",
    "            + W_l[:-1] * h_l[:-1] + W_v[:-1] * h_v[:-1]\n",
    "        )\n",
    "        return enthalpy_balance_mixture\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 实例化调用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建一个 SubchannelBlock 实例\n",
    "n_channels = 5  # 假设有5个子通道\n",
    "axial_nodes = 10  # 沿轴向有10个节点\n",
    "ndx = 9  # 通道长度的有效控制容积数\n",
    "ndxin = 9  # 输水管的有效控制容积数\n",
    "subchannel_block = SubchannelBlock(n_channels, axial_nodes, ndx, ndxin)\n",
    "\n",
    "# 初始化子通道的属性\n",
    "subchannel_block.initialize_values()\n",
    "\n",
    "# 设置一些参数\n",
    "dt = 0.01  # 时间步长\n",
    "dx = 0.1   # 空间步长\n",
    "A = 1.0    # 横截面积\n",
    "S = 0.1    # 周边面积\n",
    "K_U = 0.01 # 轴向阻力系数\n",
    "K_v = 0.05 # 横向阻力系数\n",
    "\n",
    "# 进行计算\n",
    "mass_balance = subchannel_block.mass_conservation(dt, dx)\n",
    "momentum_balance_axial = subchannel_block.axial_momentum_conservation(dt, dx, A, K_U)\n",
    "momentum_balance_transverse = subchannel_block.transverse_momentum_conservation(dt, dx, S, K_v)\n",
    "enthalpy_balance_liquid = subchannel_block.enthalpy_conservation_liquid(dt, dx, A)\n",
    "enthalpy_balance_vapor = subchannel_block.enthalpy_conservation_vapor(dt, dx, A)\n",
    "enthalpy_balance_mixture = subchannel_block.mixture_enthalpy_conservation(dt, dx, A)\n",
    "\n",
    "# 打印结果\n",
    "print(\"质量守恒方程结果:\", mass_balance)\n",
    "print(\"轴向动量守恒方程结果:\", momentum_balance_axial)\n",
    "print(\"横向动量守恒方程结果:\", momentum_balance_transverse)\n",
    "print(\"液相焓守恒方程结果:\", enthalpy_balance_liquid)\n",
    "print(\"汽相焓守恒方程结果:\", enthalpy_balance_vapor)\n",
    "print(\"混合物焓守恒方程结果:\", enthalpy_balance_mixture)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 按子通道软件进行类定义储存数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class MeshCrossSection:\n",
    "    def __init__(self, n_channels, axial_nodes):\n",
    "        self.n_channels = n_channels\n",
    "        self.axial_nodes = axial_nodes\n",
    "        self.attributes = {}  # 存储所有属性的字典\n",
    "\n",
    "        # 初始化子通道的横截面网格属性\n",
    "        self.flow_area = np.zeros((n_channels, axial_nodes))  # AREASC: 流动面积 [m²]\n",
    "        self.hydraulic_diameter = np.zeros((n_channels, axial_nodes))  # DHYD: 水力直径 [m]\n",
    "        self.heated_perimeter = np.zeros((n_channels, axial_nodes))  # HPERIM: 加热周长 [m]\n",
    "        self.roughness_height = np.zeros((n_channels, axial_nodes))  # SCHRGH: 粗糙度高度 [m]\n",
    "        self.wetted_perimeter = np.zeros((n_channels, axial_nodes))  # WPERIM: 湿润周长 [m]\n",
    "        self.inlet_drag_coefficient = np.zeros(n_channels)  # CDIN: 进口阻力系数\n",
    "        self.original_flow_area = np.zeros((n_channels, axial_nodes))  # YUANAREASC: 原始流动面积 [m²]\n",
    "        self.original_hydraulic_diameter = np.zeros((n_channels, axial_nodes))  # YUANDHYD: 原始水力直径 [m]\n",
    "        self.original_wetted_perimeter = np.zeros((n_channels, axial_nodes))  # YUANWPERIM: 原始湿润周长 [m]\n",
    "        self.hydraulic_height = np.zeros((n_channels, axial_nodes))  # HHYD: 水力高度 [m]\n",
    "        self.inlet_loss = 0  # NLOSS: 是否有进口阻力件，1表示有，0表示无\n",
    "\n",
    "    def add_attribute(self, name):\n",
    "        \"\"\"添加新的属性（如液相含率、气相速度等）到子通道块\"\"\"\n",
    "        self.attributes[name] = np.zeros((self.n_channels, self.axial_nodes + 3))\n",
    "\n",
    "    def set_attribute(self, name, values):\n",
    "        \"\"\"设置指定属性的初始值\"\"\"\n",
    "        if name in self.attributes:\n",
    "            self.attributes[name][:] = values\n",
    "        else:\n",
    "            raise ValueError(f\"Attribute {name} not defined in MeshCrossSection.\")\n",
    "\n",
    "    def get_attribute(self, name):\n",
    "        \"\"\"获取指定属性的值\"\"\"\n",
    "        return self.attributes.get(name, None)\n",
    "\n",
    "    def initialize_attributes(self):\n",
    "        \"\"\"初始化所有基础物理属性\"\"\"\n",
    "        # 示例：添加一些常见的物理属性\n",
    "        self.add_attribute(\"density\")  # 密度\n",
    "        self.add_attribute(\"flow\")  # 流量\n",
    "        self.add_attribute(\"pressure\")  # 压力\n",
    "        self.add_attribute(\"enthalpy_l\")  # 液相焓\n",
    "        self.add_attribute(\"enthalpy_g\")  # 气相焓\n",
    "        self.add_attribute(\"velocity_l\")  # 液相速度\n",
    "        self.add_attribute(\"velocity_g\")  # 气相速度\n",
    "        self.add_attribute(\"alpha_l\")  # 液相含率\n",
    "        self.add_attribute(\"alpha_g\")  # 气相含率\n",
    "class MeshLateral:\n",
    "    def __init__(self, n_gaps, axial_nodes):\n",
    "        self.n_gaps = n_gaps\n",
    "        self.axial_nodes = axial_nodes\n",
    "        self.attributes = {}  # 存储所有属性的字典\n",
    "\n",
    "        # 初始化侧向网格属性\n",
    "        self.cross_flow_area_actual = np.zeros((n_gaps, axial_nodes))  # GAPARA: 实际横向流动面积 [m²]\n",
    "        self.cross_flow_area_effective = np.zeros((n_gaps, axial_nodes))  # GAPARF: 有效横向流动面积 [m²]\n",
    "        self.centroid_distance = np.zeros((n_gaps, axial_nodes))  # GAPDST: 间隙中心距离 [m]\n",
    "        self.centroid_angle = np.zeros((n_gaps, axial_nodes))  # GAPNGL: 间隙中心角度 [弧度]\n",
    "        self.gap_width = np.zeros((n_gaps, axial_nodes))  # GAPWID: 间隙宽度 [m]\n",
    "        self.sbar = np.zeros((n_gaps, axial_nodes))  # SBAR: 修正参数，通常用于流动面积计算\n",
    "        self.cijc = np.zeros(n_gaps)  # CIJC: 修正系数\n",
    "        self.symmetry_factor = np.zeros(n_gaps)  # GAPSYM: 间隙的对称因子\n",
    "        self.rod_index = np.zeros((2, n_gaps), dtype=int)  # GAPROD: 间隙的棒段索引\n",
    "        self.ik_index = np.zeros(n_gaps, dtype=int)  # IK: 间隙中子通道 I 的索引\n",
    "        self.jk_index = np.zeros(n_gaps, dtype=int)  # JK: 间隙中子通道 J 的索引\n",
    "        self.max_difference = 0  # MAXDIF: IK 和 JK 索引的最大差值\n",
    "        self.ngap = n_gaps  # NGAP: 当前问题中的间隙数量\n",
    "\n",
    "    def add_attribute(self, name):\n",
    "        \"\"\"添加新的属性（如液相含率、气相速度等）到侧向网格\"\"\"\n",
    "        self.attributes[name] = np.zeros((self.n_gaps, self.axial_nodes + 3))\n",
    "\n",
    "    def set_attribute(self, name, values):\n",
    "        \"\"\"设置指定属性的初始值\"\"\"\n",
    "        if name in self.attributes:\n",
    "            self.attributes[name][:] = values\n",
    "        else:\n",
    "            raise ValueError(f\"Attribute {name} not defined in MeshLateral.\")\n",
    "\n",
    "    def get_attribute(self, name):\n",
    "        \"\"\"获取指定属性的值\"\"\"\n",
    "        return self.attributes.get(name, None)\n",
    "\n",
    "    def initialize_attributes(self):\n",
    "        \"\"\"初始化所有基础物理属性\"\"\"\n",
    "        # 示例：添加一些常见的物理属性\n",
    "        self.add_attribute(\"mass_exchange\")  # 质量交换\n",
    "        self.add_attribute(\"W_l_prime\")  # 液相质量交换修正\n",
    "        self.add_attribute(\"W_v_prime\")  # 气相质量交换修正\n",
    "        self.add_attribute(\"delta_U_l\")  # 液相速度修正\n",
    "        self.add_attribute(\"delta_U_v\")  # 气相速度修正\n",
    "class MeshAxial:\n",
    "    def __init__(self, n_channels, axial_nodes):\n",
    "        self.n_channels = n_channels\n",
    "        self.axial_nodes = axial_nodes\n",
    "        self.attributes = {}  # 存储所有属性的字典\n",
    "\n",
    "        # 初始化轴向网格属性\n",
    "        self.delta_x = np.zeros((n_channels, axial_nodes))  # DELX: 轴向节点长度 [m]\n",
    "        self.dx = 0  # DX: 当前轴向节点的长度 [m]\n",
    "        self.elevation = 0  # ELEV: 轴向方向的余弦值\n",
    "        self.orientation_angle = 0  # THETA: 子通道的角度 (0 度为竖直向上) [度]\n",
    "        self.x_position = np.zeros((n_channels, axial_nodes))  # X: 子通道的位置 [m]\n",
    "        self.channel_length = 0  # Z: 子通道的总长度 [m]\n",
    "        self.id_node = np.zeros((n_channels, axial_nodes), dtype=int)  # IDNODE: 轴向节点 ID\n",
    "        self.jpljm1 = 0  # JPLJM1: 轴向节点索引前一个节点\n",
    "        self.jpljp1 = 0  # JPLJP1: 轴向节点索引后一个节点\n",
    "        self.ndx = 0  # NDX: 当前问题的轴向节点数量\n",
    "        self.ndx_plus1 = 0  # NDXP1: 轴向节点数加 1\n",
    "        self.ndx_plus2 = 0  # NDXP2: 轴向节点数加 2\n",
    "\n",
    "    def add_attribute(self, name):\n",
    "        \"\"\"添加新的属性（如液相含率、气相速度等）到轴向网格\"\"\"\n",
    "        self.attributes[name] = np.zeros((self.n_channels, self.axial_nodes + 3))\n",
    "\n",
    "    def set_attribute(self, name, values):\n",
    "        \"\"\"设置指定属性的初始值\"\"\"\n",
    "        if name in self.attributes:\n",
    "            self.attributes[name][:] = values\n",
    "        else:\n",
    "            raise ValueError(f\"Attribute {name} not defined in MeshAxial.\")\n",
    "\n",
    "    def get_attribute(self, name):\n",
    "        \"\"\"获取指定属性的值\"\"\"\n",
    "        return self.attributes.get(name, None)\n",
    "\n",
    "    def initialize_attributes(self):\n",
    "        \"\"\"初始化所有基础物理属性\"\"\"\n",
    "        # 示例：添加一些常见的物理属性\n",
    "        self.add_attribute(\"gravity_x\")  # 重力在x方向的加速度\n",
    "        self.add_attribute(\"gravity_y\")  # 重力在y方向的加速度\n",
    "        \n",
    "def conjugate_gradient(A, b, x0=None, tol=1e-6, max_iter=None):\n",
    "    \"\"\"共轭梯度法求解线性方程组 Ax = b\"\"\"\n",
    "    # 确保输入为numpy数组\n",
    "    A = np.array(A)\n",
    "    b = np.array(b)\n",
    "\n",
    "    # 如果未指定初始解，使用全零向量\n",
    "    if x0 is None:\n",
    "        x0 = np.zeros_like(b)\n",
    "\n",
    "    # 如果未指定最大迭代次数，使用矩阵维度\n",
    "    if max_iter is None:\n",
    "        max_iter = len(b)\n",
    "\n",
    "    # 初始化\n",
    "    x = x0\n",
    "    r = b - A @ x  # 残差\n",
    "    p = r.copy()   # 搜索方向\n",
    "\n",
    "    # 迭代求解\n",
    "    for _ in range(max_iter):\n",
    "        # 计算步长\n",
    "        alpha = np.dot(r, r) / np.dot(p, A @ p)\n",
    "\n",
    "        # 更新解和残差\n",
    "        x = x + alpha * p\n",
    "        r_next = r - alpha * (A @ p)\n",
    "\n",
    "        # 检查收敛性\n",
    "        if np.linalg.norm(r_next) < tol:\n",
    "            return x\n",
    "\n",
    "        # 计算beta（共轭方向系数）\n",
    "        beta = np.dot(r_next, r_next) / np.dot(r, r)\n",
    "\n",
    "        # 更新搜索方向\n",
    "        p = r_next + beta * p\n",
    "\n",
    "        # 更新残差\n",
    "        r = r_next\n",
    "\n",
    "    return x\n",
    "\n",
    "class ConservationEquations:\n",
    "    def __init__(self, cross_section, lateral, axial):\n",
    "        self.cross_section = cross_section\n",
    "        self.lateral = lateral\n",
    "        self.axial = axial\n",
    "\n",
    "    def mass_conservation(self, dt, dx):\n",
    "        rho = self.cross_section.get_attribute(\"density\")\n",
    "        F = self.cross_section.get_attribute(\"flow\")\n",
    "        W = self.lateral.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        drho_dt = (rho[1:] - rho[:-1]) / dt\n",
    "        dF_dx = (F[1:] - F[:-1]) / dx\n",
    "        mass_balance = drho_dt + dF_dx + W[:-1]\n",
    "        return mass_balance\n",
    "    \n",
    "    def axial_momentum_conservation(self, dt, dx, A, K_U):\n",
    "        F = self.cross_section.get_attribute(\"flow\")\n",
    "        P = self.cross_section.get_attribute(\"pressure\")\n",
    "        rho = self.cross_section.get_attribute(\"density\")\n",
    "        g_x = self.axial.get_attribute(\"gravity_x\")\n",
    "        \n",
    "        F_l = self.cross_section.get_attribute(\"flow_l\")\n",
    "        F_v = self.cross_section.get_attribute(\"flow_v\")\n",
    "        U_l = self.cross_section.get_attribute(\"velocity_l\")\n",
    "        U_v = self.cross_section.get_attribute(\"velocity_g\")\n",
    "        \n",
    "        W_l = self.lateral.get_attribute(\"mass_exchange\")\n",
    "        W_l_prime = self.lateral.get_attribute(\"W_l_prime\")\n",
    "        delta_U_l = self.lateral.get_attribute(\"delta_U_l\")\n",
    "        W_v_prime = self.lateral.get_attribute(\"W_v_prime\")\n",
    "        delta_U_v = self.lateral.get_attribute(\"delta_U_v\")\n",
    "        \n",
    "        dF_dt = (F[1:] - F[:-1]) / dt\n",
    "        dF_dx = (F_l[1:] * U_l[1:] + F_v[1:] * U_v[1:] - F_l[:-1] * U_l[:-1] - F_v[:-1] * U_v[:-1]) / dx\n",
    "        dP_dx = (P[1:] - P[:-1]) / dx\n",
    "        \n",
    "        momentum_balance = (\n",
    "            dF_dt + dF_dx\n",
    "            + dP_dx * A\n",
    "            - W_l[:-1] * U_l[:-1]\n",
    "            - W_v_prime[:-1] * delta_U_v[:-1]\n",
    "            - K_U * np.abs(F[:-1]) * F[:-1]\n",
    "            + rho[:-1] * g_x[:-1]\n",
    "        )\n",
    "        return momentum_balance\n",
    "    \n",
    "    def transverse_momentum_conservation(self, dt, dx, S, K_v):\n",
    "        rho = self.cross_section.get_attribute(\"density\")\n",
    "        V = self.cross_section.get_attribute(\"velocity_g\")\n",
    "        F = self.cross_section.get_attribute(\"flow\")\n",
    "        P = self.cross_section.get_attribute(\"pressure\")\n",
    "        W = self.lateral.get_attribute(\"mass_exchange\")\n",
    "        W_v_prime = self.lateral.get_attribute(\"W_v_prime\")\n",
    "        delta_U_v = self.lateral.get_attribute(\"delta_U_v\")\n",
    "        g_y = 9.81  # 重力在y方向的加速度\n",
    "        \n",
    "        dV_dt = (V[1:] - V[:-1]) / dt\n",
    "        dFV_dx = (F[1:] * V[1:] - F[:-1] * V[:-1]) / dx\n",
    "        dF_dx = (F[1:] - F[:-1]) / dx\n",
    "        dP_dy = (P[1:] - P[:-1]) / dx  # 近似在x方向\n",
    "        \n",
    "        transverse_balance = (\n",
    "            9 * rho[:-1] * dV_dt + dFV_dx - V[:-1] * dF_dx\n",
    "            - 9 * dP_dy\n",
    "            + S * dx * K_v * W[:-1] * W[:-1]\n",
    "            + (W[:-1] * delta_U_v[:-1] + W_v_prime[:-1] * delta_U_v[:-1]) / (S * dx)\n",
    "            + 9 * rho[:-1] * g_y\n",
    "        )\n",
    "        return transverse_balance\n",
    "    \n",
    "    def enthalpy_conservation_liquid(self, dt, dx, A):\n",
    "        alpha = self.cross_section.get_attribute(\"alpha_l\")\n",
    "        rho_l = self.cross_section.get_attribute(\"density_l\")\n",
    "        h_l = self.cross_section.get_attribute(\"enthalpy_l\")\n",
    "        F_l = self.cross_section.get_attribute(\"flow_l\")\n",
    "        W_l = self.lateral.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        dh_l_dt = (h_l[1:] - h_l[:-1]) / dt\n",
    "        dFh_dx = (F_l[1:] * h_l[1:] - F_l[:-1] * h_l[:-1]) / dx\n",
    "        dF_dx = (F_l[1:] - F_l[:-1]) / dx\n",
    "        \n",
    "        enthalpy_balance_l = (\n",
    "            (1 - alpha[:-1]) * A * dx * rho_l[:-1] * dh_l_dt\n",
    "            + dx * dFh_dx - dx * h_l[:-1] * dF_dx + W_l[:-1] * h_l[:-1]\n",
    "        )\n",
    "        return enthalpy_balance_l\n",
    "    \n",
    "    def enthalpy_conservation_vapor(self, dt, dx, A):\n",
    "        alpha = self.cross_section.get_attribute(\"alpha_g\")\n",
    "        rho_v = self.cross_section.get_attribute(\"density_v\")\n",
    "        h_v = self.cross_section.get_attribute(\"enthalpy_g\")\n",
    "        F_v = self.cross_section.get_attribute(\"flow_v\")\n",
    "        W_v = self.lateral.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        dh_v_dt = (h_v[1:] - h_v[:-1]) / dt\n",
    "        dFh_dx = (F_v[1:] * h_v[1:] - F_v[:-1] * h_v[:-1]) / dx\n",
    "        dF_dx = (F_v[1:] - F_v[:-1]) / dx\n",
    "        \n",
    "        enthalpy_balance_v = (\n",
    "            alpha[:-1] * A * dx * rho_v[:-1] * dh_v_dt\n",
    "            + dx * dFh_dx - dx * h_v[:-1] * dF_dx + W_v[:-1] * h_v[:-1]\n",
    "        )\n",
    "        return enthalpy_balance_v\n",
    "    \n",
    "    def mixture_enthalpy_conservation(self, dt, dx, A):\n",
    "        rho = self.cross_section.get_attribute(\"density\")\n",
    "        h = self.cross_section.get_attribute(\"enthalpy_l\")  # 使用液相焓作为混合物焓\n",
    "        F_l = self.cross_section.get_attribute(\"flow_l\")\n",
    "        h_l = self.cross_section.get_attribute(\"enthalpy_l\")\n",
    "        F_v = self.cross_section.get_attribute(\"flow_v\")\n",
    "        h_v = self.cross_section.get_attribute(\"enthalpy_g\")\n",
    "        W_l = self.lateral.get_attribute(\"mass_exchange\")\n",
    "        W_v = self.lateral.get_attribute(\"mass_exchange\")\n",
    "        \n",
    "        dh_dt = (h[1:] - h[:-1]) / dt\n",
    "        dFh_dx = (F_l[1:] * h_l[1:] + F_v[1:] * h_v[1:] - F_l[:-1] * h_l[:-1] - F_v[:-1] * h_v[:-1]) / dx\n",
    "        dF_dx = (F_l[1:] + F_v[1:] - F_l[:-1] - F_v[:-1]) / dx\n",
    "        \n",
    "        enthalpy_balance_mixture = (\n",
    "            A * dx * rho[:-1] * dh_dt\n",
    "            + dx * dFh_dx - dx * h[:-1] * dF_dx\n",
    "            + W_l[:-1] * h_l[:-1] + W_v[:-1] * h_v[:-1]\n",
    "        )\n",
    "        return enthalpy_balance_mixture\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.42851054 0.91469071 0.02437832 ... 0.80188488 0.85459466 0.62162211]\n",
      " [0.17130994 0.78695852 0.2426596  ... 0.19302753 0.13890519 0.07427075]\n",
      " [0.31201547 0.42459771 0.89103089 ... 0.64283417 0.78957729 0.24015745]\n",
      " ...\n",
      " [0.8669784  0.75440993 0.3011151  ... 0.4637786  0.38378783 0.34699147]\n",
      " [0.49837587 0.51696304 0.04191715 ... 0.03343867 0.0565777  0.50985089]\n",
      " [0.5548927  0.12367303 0.10841695 ... 0.59096449 0.77219291 0.93545562]]\n",
      "[[0.18287677 0.72368674 0.88901453 0.70185613 0.90601343 0.24636775\n",
      "  0.85274376 0.74387345 0.00221122 0.94886613 0.91214335 0.17934294\n",
      "  0.24280457 0.79104133 0.27646129 0.65032056 0.14056621 0.27227691\n",
      "  0.91409744 0.05253896 0.63295596 0.50369115 0.76120205 0.934699\n",
      "  0.3762481  0.16502932 0.0843462  0.6025754  0.69762574 0.31999569\n",
      "  0.92978817 0.77192953 0.68985198 0.90951393 0.13257556 0.17923781\n",
      "  0.85504662 0.61901225 0.12954499 0.11028562 0.1435795  0.53695211\n",
      "  0.30474845 0.18735319 0.54975035 0.14887512 0.38220893 0.44346908\n",
      "  0.17724453 0.147158   0.07335391 0.45135763 0.2431813 ]\n",
      " [0.48346713 0.95723442 0.89020841 0.18370024 0.09358864 0.23503227\n",
      "  0.79272396 0.3853698  0.48004192 0.62886216 0.88864144 0.97537252\n",
      "  0.72603457 0.18152164 0.76152713 0.03607887 0.87306482 0.00713483\n",
      "  0.87227999 0.07856516 0.41910379 0.2996459  0.12942126 0.39712762\n",
      "  0.46490816 0.71895847 0.8613483  0.71703264 0.00713708 0.81386928\n",
      "  0.01608409 0.33754055 0.6366385  0.95426435 0.98895636 0.12197061\n",
      "  0.58077419 0.49246436 0.25996265 0.27558196 0.17569616 0.57329791\n",
      "  0.80241892 0.47343088 0.46238481 0.18131668 0.98845362 0.34898069\n",
      "  0.05966264 0.62246792 0.24360192 0.27643035 0.7917005 ]\n",
      " [0.99032002 0.6070791  0.13208137 0.89531802 0.08440654 0.27437494\n",
      "  0.69872945 0.56026138 0.57100087 0.98790712 0.3792219  0.22062079\n",
      "  0.20364122 0.23396346 0.96755945 0.04825281 0.87274007 0.3102111\n",
      "  0.935154   0.3134141  0.87843629 0.12803195 0.4360254  0.68638922\n",
      "  0.15434583 0.9559835  0.89794013 0.18811665 0.00473426 0.12054864\n",
      "  0.26631234 0.07678364 0.47526155 0.79110098 0.69723967 0.75360872\n",
      "  0.38233104 0.74285818 0.23898037 0.8882102  0.73916652 0.97312351\n",
      "  0.60445697 0.93201123 0.2550904  0.62375077 0.60700867 0.41943399\n",
      "  0.39376511 0.18514746 0.93231991 0.41917316 0.67891771]\n",
      " [0.87803149 0.57812513 0.82806949 0.6914385  0.39092778 0.71805172\n",
      "  0.08724531 0.78904766 0.78875995 0.38539194 0.911415   0.73352572\n",
      "  0.41536805 0.53967422 0.25299314 0.61108734 0.63146005 0.64928361\n",
      "  0.79341872 0.94894339 0.51398423 0.67585406 0.45314941 0.57865692\n",
      "  0.54721086 0.21789831 0.94686945 0.51535287 0.34338606 0.51694072\n",
      "  0.51242415 0.30931456 0.10472203 0.23143465 0.84899864 0.83416249\n",
      "  0.49007186 0.22232972 0.18127137 0.79631778 0.59075122 0.99035942\n",
      "  0.7337519  0.46727072 0.37347151 0.91209601 0.94872215 0.38678633\n",
      "  0.08472047 0.6777221  0.70002603 0.22575126 0.77292761]\n",
      " [0.23022202 0.0604326  0.12909575 0.92449585 0.3273425  0.64028302\n",
      "  0.42132526 0.79103385 0.10954237 0.6334651  0.47889772 0.40061612\n",
      "  0.48013377 0.02313327 0.36689361 0.04044101 0.94807577 0.91431975\n",
      "  0.26247969 0.25735478 0.6127299  0.81188943 0.03677108 0.45773081\n",
      "  0.12222628 0.39536364 0.06724386 0.98638467 0.09839785 0.08383376\n",
      "  0.23822898 0.13442039 0.11437463 0.94065621 0.46923574 0.82534112\n",
      "  0.23043587 0.74742788 0.97317025 0.79910887 0.74463541 0.65130218\n",
      "  0.1346043  0.07487145 0.50000627 0.86333208 0.36241583 0.71379024\n",
      "  0.29593068 0.44554075 0.63766316 0.71851819 0.93653918]\n",
      " [0.35858377 0.31649074 0.21230338 0.57607687 0.10671963 0.24114876\n",
      "  0.57709785 0.76226865 0.32760478 0.89963672 0.81726292 0.44513454\n",
      "  0.73583664 0.58875123 0.64759286 0.54165032 0.1690762  0.60898231\n",
      "  0.34601658 0.09746934 0.05671067 0.75405882 0.96577525 0.95955154\n",
      "  0.10729353 0.39169581 0.21734745 0.14102453 0.74248868 0.73674561\n",
      "  0.04611    0.61078749 0.49084579 0.7717838  0.7214139  0.1750584\n",
      "  0.8287535  0.8570456  0.32996101 0.22011739 0.01579564 0.03669361\n",
      "  0.32635275 0.01656424 0.07810827 0.62659286 0.22077811 0.80211783\n",
      "  0.12908853 0.62516129 0.95281182 0.73324613 0.2722275 ]\n",
      " [0.46619624 0.71573966 0.99842115 0.22079941 0.15636938 0.55170213\n",
      "  0.93882859 0.83686543 0.60194471 0.22395051 0.49550428 0.15800667\n",
      "  0.50800185 0.06066883 0.15036657 0.94349909 0.21746686 0.78661641\n",
      "  0.58517838 0.27635304 0.24355296 0.58859758 0.56010674 0.65921927\n",
      "  0.22154929 0.88305152 0.95541422 0.76931203 0.50166873 0.70664832\n",
      "  0.28109163 0.33761798 0.73665109 0.03890677 0.5648511  0.47574437\n",
      "  0.14940317 0.43613286 0.79400497 0.16882988 0.62932941 0.93752084\n",
      "  0.39108074 0.04763078 0.76987522 0.13740656 0.28771156 0.11065088\n",
      "  0.9224626  0.09003961 0.54719888 0.62491821 0.64810224]\n",
      " [0.03502601 0.57891795 0.94076733 0.04740759 0.61559634 0.52299217\n",
      "  0.46299649 0.03086616 0.01405211 0.96527635 0.97379107 0.64378578\n",
      "  0.49491231 0.23203204 0.82178664 0.53974987 0.43991641 0.09119304\n",
      "  0.23922719 0.9050575  0.05822178 0.47900352 0.66559373 0.58368183\n",
      "  0.9315893  0.0016217  0.39947997 0.63889232 0.056925   0.27611347\n",
      "  0.49708795 0.59669719 0.45491997 0.73146755 0.5694996  0.36964179\n",
      "  0.98834236 0.48948714 0.35448572 0.0329656  0.43262515 0.91029366\n",
      "  0.38801028 0.44735742 0.4383264  0.43555809 0.55099872 0.17711325\n",
      "  0.2558956  0.51342324 0.78446442 0.49777279 0.07922658]\n",
      " [0.19399336 0.53507688 0.39745926 0.87126898 0.56904592 0.571\n",
      "  0.59992653 0.05972111 0.79191832 0.31650612 0.06700208 0.07492932\n",
      "  0.61871253 0.92095833 0.74946051 0.84682934 0.59007249 0.42345384\n",
      "  0.90685652 0.09852173 0.15319269 0.47158385 0.42193184 0.03876324\n",
      "  0.05148266 0.80610372 0.90882744 0.47265913 0.63310494 0.97115354\n",
      "  0.16463178 0.52547628 0.7723481  0.8292431  0.16614416 0.11931557\n",
      "  0.46589674 0.70739226 0.14308993 0.14252498 0.26754019 0.60611768\n",
      "  0.75917107 0.31616751 0.66594992 0.87701096 0.52961879 0.31623641\n",
      "  0.36536008 0.13531286 0.52972536 0.16271413 0.04708577]\n",
      " [0.82481393 0.15478846 0.05171424 0.25448094 0.59242357 0.99706526\n",
      "  0.54359035 0.58469199 0.13906271 0.60647343 0.46221222 0.44763241\n",
      "  0.52245691 0.14372819 0.85403844 0.42363244 0.52587618 0.2037118\n",
      "  0.0475705  0.4427018  0.35507426 0.68047694 0.38358268 0.46850737\n",
      "  0.86251774 0.04502854 0.37334809 0.43990183 0.61429285 0.07892806\n",
      "  0.78012637 0.76211097 0.87084902 0.42139658 0.87368066 0.79018889\n",
      "  0.35851853 0.24204884 0.77721551 0.22980521 0.98991345 0.9354868\n",
      "  0.92487724 0.1804865  0.16747883 0.73419153 0.48296779 0.78336161\n",
      "  0.65802233 0.27720791 0.38247525 0.86569058 0.98617808]]\n",
      "[[0.46984715 0.56970593 0.85972759 ... 0.42619056 0.37801063 0.63988528]\n",
      " [0.75840304 0.76363235 0.70902972 ... 0.94536968 0.06495944 0.28030479]\n",
      " [0.31191867 0.31100823 0.36466367 ... 0.67182561 0.80881329 0.71735709]\n",
      " ...\n",
      " [0.54003985 0.83288676 0.0415808  ... 0.13471965 0.98337078 0.57212777]\n",
      " [0.28849464 0.18332529 0.82794134 ... 0.78090944 0.78904876 0.21381618]\n",
      " [0.04721984 0.03233581 0.29181445 ... 0.84654865 0.45294161 0.08187224]]\n"
     ]
    }
   ],
   "source": [
    "cross_section = MeshCrossSection(n_channels=70, axial_nodes=50)\n",
    "lateral = MeshLateral(n_gaps=10, axial_nodes=50)\n",
    "axial = MeshAxial(n_channels=70, axial_nodes=50)\n",
    "\n",
    "# 初始化属性\n",
    "cross_section.initialize_attributes()\n",
    "lateral.initialize_attributes()\n",
    "axial.initialize_attributes()\n",
    "\n",
    "# 设置属性值\n",
    "cross_section.set_attribute(\"density\", np.random.rand(70, 53))\n",
    "lateral.set_attribute(\"mass_exchange\", np.random.rand(10, 53))\n",
    "axial.set_attribute(\"gravity_x\", np.random.rand(70, 53))\n",
    "\n",
    "# 获取属性值\n",
    "density = cross_section.get_attribute(\"density\")\n",
    "mass_exchange = lateral.get_attribute(\"mass_exchange\")\n",
    "gravity_x = axial.get_attribute(\"gravity_x\")\n",
    "\n",
    "print(density)\n",
    "print(mass_exchange)\n",
    "print(gravity_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 在调用守恒方程时报错了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "operands could not be broadcast together with shapes (69,53) (9,53) ",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_16240\\1138675626.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[0mdt\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m0.1\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     23\u001b[0m \u001b[0mdx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m0.01\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 24\u001b[1;33m \u001b[0mmass_balance\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mconservation_eqs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmass_conservation\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdt\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     25\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     26\u001b[0m \u001b[1;31m# 计算轴向动量守恒\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\AppData\\Local\\Temp\\ipykernel_16240\\2382486561.py\u001b[0m in \u001b[0;36mmass_conservation\u001b[1;34m(self, dt, dx)\u001b[0m\n\u001b[0;32m     12\u001b[0m         \u001b[0mdrho_dt\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mrho\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m-\u001b[0m \u001b[0mrho\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m/\u001b[0m \u001b[0mdt\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     13\u001b[0m         \u001b[0mdF_dx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mF\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m-\u001b[0m \u001b[0mF\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m/\u001b[0m \u001b[0mdx\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 14\u001b[1;33m         \u001b[0mmass_balance\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mdrho_dt\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mdF_dx\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mW\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     15\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0mmass_balance\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     16\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mValueError\u001b[0m: operands could not be broadcast together with shapes (69,53) (9,53) "
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# 初始化网格\n",
    "cross_section = MeshCrossSection(n_channels=70, axial_nodes=50)\n",
    "lateral = MeshLateral(n_gaps=10, axial_nodes=50)\n",
    "axial = MeshAxial(n_channels=70, axial_nodes=50)\n",
    "\n",
    "# 初始化属性\n",
    "cross_section.initialize_attributes()\n",
    "lateral.initialize_attributes()\n",
    "axial.initialize_attributes()\n",
    "\n",
    "# 设置属性值\n",
    "cross_section.set_attribute(\"density\", np.random.rand(70, 53))\n",
    "lateral.set_attribute(\"mass_exchange\", np.random.rand(10, 53))\n",
    "axial.set_attribute(\"gravity_x\", np.random.rand(70, 53))\n",
    "\n",
    "# 初始化守恒方程\n",
    "conservation_eqs = ConservationEquations(cross_section, lateral, axial)\n",
    "\n",
    "# 计算质量守恒\n",
    "dt = 0.1\n",
    "dx = 0.01\n",
    "mass_balance = conservation_eqs.mass_conservation(dt, dx)\n",
    "\n",
    "# 计算轴向动量守恒\n",
    "A = 1.0\n",
    "K_U = 0.5\n",
    "axial_momentum_balance = conservation_eqs.axial_momentum_conservation(dt, dx, A, K_U)\n",
    "\n",
    "# 计算横向动量守恒\n",
    "S = 1.0\n",
    "K_v = 0.5\n",
    "transverse_momentum_balance = conservation_eqs.transverse_momentum_conservation(dt, dx, S, K_v)\n",
    "\n",
    "# 计算液相焓守恒\n",
    "enthalpy_balance_liquid = conservation_eqs.enthalpy_conservation_liquid(dt, dx, A)\n",
    "\n",
    "# 计算气相焓守恒\n",
    "enthalpy_balance_vapor = conservation_eqs.enthalpy_conservation_vapor(dt, dx, A)\n",
    "\n",
    "# 计算混合物焓守恒\n",
    "enthalpy_balance_mixture = conservation_eqs.mixture_enthalpy_conservation(dt, dx, A)\n",
    "\n",
    "# 打印结果\n",
    "print(\"质量守恒方程结果:\", mass_balance)\n",
    "print(\"轴向动量守恒方程结果:\", axial_momentum_balance)\n",
    "print(\"横向动量守恒方程结果:\", transverse_momentum_balance)\n",
    "print(\"液相焓守恒方程结果:\", enthalpy_balance_liquid)\n",
    "print(\"汽相焓守恒方程结果:\", enthalpy_balance_vapor)\n",
    "print(\"混合物焓守恒方程结果:\", enthalpy_balance_mixture)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
